import os
import cv2
import numpy as np


def normalize_images(dataset_path, output_path):
    """将数据集中所有图像归一化为640x640大小，保持原始纵横比并使用填充"""
    os.makedirs(output_path, exist_ok=True)

    target_size = 640  # 在此处修改大小

    for filename in os.listdir(dataset_path):
        image_path = os.path.join(dataset_path, filename)
        image = cv2.imread(image_path)

        if image is not None:
            # 获取原始尺寸
            h, w = image.shape[:2]
            # 计算缩放比例: 通过把640除以h或w，获得最小的那个作为缩放比例
            scale = min(target_size / h, target_size / w)
            # 计算新的尺寸
            new_w = int(w * scale)
            new_h = int(h * scale)

            # Resize图像以保持纵横比
            resized_image = cv2.resize(image, (new_w, new_h))

            # 创建填充后的图像，大小为640x640，目前为全黑
            padded_image = np.zeros((target_size, target_size, 3), dtype=np.uint8)  # 黑色背景
            # 计算填充位置
            # 将调整大小后的图像放置到填充图像居中的位置
            x_offset = (target_size - new_w) // 2
            y_offset = (target_size - new_h) // 2

            # 将图像放入填充图像中
            padded_image[y_offset:y_offset + new_h, x_offset:x_offset + new_w] = resized_image

            # 保存归一化后的图像
            output_file_path = os.path.join(output_path, filename)
            cv2.imwrite(output_file_path, padded_image)
            print(f"Normalized image saved as: {output_file_path}")


if __name__ == "__main__":
    # 输入和输出文件夹路径
    dataset_path = 'VOC2028_short/images1'  # 替换为原始数据集路径
    output_path = 'VOC2028_short/images4'  # 替换为输出路径

    # 调用函数归一化图片
    normalize_images(dataset_path, output_path)